from ultralytics import YOLO
import cv2

# weight_path = r"E:\项目\松山湖公安分局无人机自动巡检项目\事故检测\标注数据\OBB\权重文件\模型精度提升训练_样本数4586（目前较好）\train62\weights\best.pt"
# weight_path = r"E:\项目\松山湖公安分局无人机自动巡检项目\红外异常活动识别\数据集\yolov5_trained_model\yolov5_trained_model\yolov5_infrared.pt"
# weight_path = r"C:\Users\14159\Desktop\其他工具\工具箱\ai\runs\detect\train13\weights\best.pt"
# weight_path = r"E:\项目\松山湖公安分局无人机自动巡检项目\事故检测\标注数据\OBB\交通事故0716\权重\train36\weights\best.pt"
# weight_path = r"D:\code_work\git\model_library\weight\yolo_accident_obb.pt"
# video_path = r"rtmp://video.godouav.com/live/1816437285179953152181735212?appId=zkytApp001"
# video_path = r"rtmp://10.5.52.56:1935/live/1581F6Q8D248E00G0M54?callId=26-173"
# video_path = r"rtmp://10.1.38.245:1935/live/raw_stream3"
# video_path = r"E:\项目\松山湖公安分局无人机自动巡检项目\红外异常活动识别\无人机数据-云遥\250325\红外识别\1816437285179953152181735212_1742814728-00.04.20.563-00.10.16.453-seg2.mov"
# video_path = r"E:\项目\松山湖公安分局无人机自动巡检项目\事故检测\无人机事故视频-云遥\模型能力提升阶段一\第四周\交通事故0811\1816437285179953152181735212_20250811110708-00.07.35.937-00.14.04.244-seg2.mp4"

weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\public_security\traffic_accient\20250925\train77\weights\best.pt"
video_path = r"E:\zhihao\2025文档\25年10月\图片\1.mp4"


# 打开视频文件
cap = cv2.VideoCapture(video_path,cv2.CAP_FFMPEG)
# 获取视频的基本信息
fps = int(cap.get(cv2.CAP_PROP_FPS))  # 获取原视频帧率
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 视频宽度
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 视频高度
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 总帧数
print(fps,width,height,frame_count)

model = YOLO(weight_path)

# model.track(video_path, imgsz=(height,width),classes=[0], rect=True, iou=0.3, conf=0.6,visualize=False,vid_stride=2,half=True, save=False,show=False)
# model.predict(video_path,vid_stride=2,classes=[3,4,5], save=True,show=False,show_labels=True,show_conf=False)
#通用
model.predict(video_path,vid_stride=1,imgsz=(height,width),classes=[0],save=True,show=True,show_labels=True,show_conf=True,conf=0.5,iou=0.3)


